package main

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/pem"
	"github.com/emmansun/gmsm/pkcs8"
	"github.com/emmansun/gmsm/sm2"
	"github.com/emmansun/gmsm/smx509"
	"io"
	"log"
	"os"
)

func GenerateKey() *ecdsa.PrivateKey {
	priKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		log.Println(err)
	}

	return &priKey.PrivateKey
}

func GenerateKeyWithFile() (*ecdsa.PrivateKey, *ecdsa.PublicKey) {
	priKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		log.Println(err)
	}
	derStream, err := pkcs8.MarshalPrivateKey(priKey, []byte(""), pkcs8.DefaultOpts)
	//derStream, err := x509.MarshalPKCS8PrivateKey(&priKey.PrivateKey)
	//smx509.MarshalSM2PrivateKey()
	if err != nil {
		panic(err)
	}
	block := pem.Block{
		Type:  "SM2 PRIVATE KEY", //"RSA PRIVATE KEY",
		Bytes: derStream,
	}
	// 4. 创建文件
	privFile, err := os.Create("./key/private.pem")
	if err != nil {
		panic(err)
	}
	defer privFile.Close()
	// 5. 使用pem编码, 并将数据写入文件中
	err = pem.Encode(privFile, &block)
	if err != nil {
		panic(err)
	}
	// 6. 最后的时候关闭文件
	defer privFile.Close()

	//生成公钥文件
	pubKey := priKey.Public().(*ecdsa.PublicKey)
	derPkix, err := smx509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		panic(err)
	}
	block = pem.Block{
		Type:  "SM2 PUBLIC KEY", //"PUBLIC KEY",
		Bytes: derPkix,
	}
	pubFile, err := os.Create("./key/public.pem")
	defer pubFile.Close()
	if err != nil {
		panic(err)
	}
	// 8. 编码公钥, 写入文件
	err = pem.Encode(pubFile, &block)
	if err != nil {
		panic(err)
	}
	defer pubFile.Close()

	return &priKey.PrivateKey, pubKey
}

func sm2ToBytes(pr *ecdsa.PrivateKey) (prbytes []byte, pubytes []byte) {
	b1 := new(bytes.Buffer)
	derStream, err := pkcs8.MarshalPrivateKey(ECDSATOSM2(pr), []byte(""), pkcs8.DefaultOpts)
	if err != nil {
		panic(err)
	}
	block := pem.Block{
		Type:  "SM2 PRIVATE KEY", //"RSA PRIVATE KEY",
		Bytes: derStream,
	}
	err = pem.Encode(b1, &block)
	if err != nil {
		panic(err)
	}

	b2 := new(bytes.Buffer)
	pubKey := pr.Public().(*ecdsa.PublicKey)
	derPkix, err := smx509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		panic(err)
	}
	block = pem.Block{
		Type:  "SM2 PUBLIC KEY", //"PUBLIC KEY",
		Bytes: derPkix,
	}
	pem.Encode(b2, &block)
	if err != nil {
		panic(err)
	}

	//fmt.Println("b1.Bytes(), b2.Bytes()", string(b1.Bytes()), string(b1.Bytes()))
	return b1.Bytes(), b2.Bytes()
}

// ECDSATOSM2 将ecdsa.PrivateKey转为sm2.PrivateKey
func ECDSATOSM2(pr *ecdsa.PrivateKey) *sm2.PrivateKey {
	privateKey, err := sm2.NewPrivateKey(pr.D.Bytes())
	if err != nil {
		panic(err)
	}
	return privateKey
}

func ReadKeyFromFile() (pr *ecdsa.PrivateKey, pu *ecdsa.PublicKey) {
	//pubFile := "./key/public.pem"
	//priFile := "./key/private.pem"

	pubF, err := os.Open("./key/public.pem")
	priF, err := os.Open("./key/private.pem")
	defer pubF.Close()
	defer priF.Close()
	if err != nil {
		panic(err)
	}
	//解析私钥
	b1 := new(bytes.Buffer)
	io.Copy(b1, priF)
	decode1, _ := pem.Decode(b1.Bytes())
	prikey, err := smx509.ParsePKCS8PrivateKey(decode1.Bytes)
	if err != nil {
		panic(err)
	}
	//注意pem读文件后是*sm2.PrivateKey格式
	pr = &prikey.(*sm2.PrivateKey).PrivateKey

	//解析公钥
	b2 := new(bytes.Buffer)
	io.Copy(b2, pubF)
	decode2, _ := pem.Decode(b2.Bytes())
	key, err := smx509.ParsePKIXPublicKey(decode2.Bytes)
	if err != nil {
		panic(err)
	}
	pu = key.(*ecdsa.PublicKey)

	return
}
